package com.googlecode.objectify.gwtkey;

import static javax.lang.model.SourceVersion.RELEASE_6;

import java.io.BufferedWriter;
import java.io.PrintWriter;
import java.util.Set;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.annotation.processing.SupportedAnnotationTypes;
import javax.annotation.processing.SupportedSourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import javax.lang.model.element.VariableElement;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.ElementFilter;
import javax.tools.Diagnostic.Kind;

import com.googlecode.objectify.annotation.Entity;
import javax.persistence.Id;
import com.googlecode.objectify.annotation.Parent;
import com.googlecode.objectify.annotation.Subclass;

@SupportedAnnotationTypes("*")
@SupportedSourceVersion(RELEASE_6)
public class EntityProcessor extends AbstractProcessor {

  private ProcessingEnvironment env;

  @Override
  public synchronized void init(ProcessingEnvironment env) {
    super.init(env);

    this.env = env;
  }

  @Override
  public boolean process(Set<? extends TypeElement> annotations,
      RoundEnvironment roundEnv) {

    if (!roundEnv.processingOver()) {
  
      this.env.getMessager().printMessage(Kind.NOTE, "com.googlecode.objectify.gwtkey.EntityProcessor started.");

      this.env.getMessager().printMessage(Kind.NOTE, "Searching for @Entity annotations.");
      for (Element entity : roundEnv.getElementsAnnotatedWith(Entity.class)) {
        this.env.getMessager().printMessage(Kind.NOTE, "Found " + entity.toString() + ".");
        if(!this.generateKeyFactory((TypeElement) entity)) {
          return false;
        }
      }
    
      this.env.getMessager().printMessage(Kind.NOTE, "Searching for @Subclass annotations.");
      for (Element subclass : roundEnv.getElementsAnnotatedWith(Subclass.class)) {
        this.env.getMessager().printMessage(Kind.NOTE, "Found @" + subclass.toString() + ".");
        if(!this.generateKeyFactory((TypeElement) subclass)) {
          return false;
        }
      }
    
      this.env.getMessager().printMessage(Kind.NOTE, "com.googlecode.objectify.gwtkey.EntityProcessor finished.");
    }

    return true;
  }

  private boolean generateKeyFactory(TypeElement entityElement) {

    PrintWriter out = null;

    String entityPackageName = env.getElementUtils().getPackageOf(entityElement).getQualifiedName().toString();
    String entityName = entityElement.getSimpleName().toString();
    String keyFactoryName = entityName + "KeyFactory";
    String keyFactoryQName = entityPackageName + "." + keyFactoryName;

    this.env.getMessager().printMessage(Kind.NOTE, "Generating '" + keyFactoryQName + "' from '" + entityName + "'.");

    Entity entityAnnotation = entityElement.getAnnotation(Entity.class);
    Subclass subclassAnnotation = entityElement.getAnnotation(Subclass.class);
    String kind = null;
    if (entityAnnotation != null) {
      if(entityAnnotation.name().isEmpty()) {
        kind = entityElement.getSimpleName().toString();
      } else {
        kind = entityAnnotation.name();
      }
    } else if (subclassAnnotation != null) {
      TypeMirror superClass = entityElement.getSuperclass();
      Entity superClassEntityAnnotation = null;
      TypeElement superClassElement = null;
      do {
        if (superClass.getKind() == TypeKind.NONE) {
          this.env.getMessager().printMessage(
              Kind.NOTE,
              "ERROR: @Subclass " + entityName + " did not have superclass with @Entity");
          return false;
        }
        superClassElement = (TypeElement) env.getTypeUtils().asElement(
            superClass);
        superClassEntityAnnotation = superClassElement.getAnnotation(Entity.class);
      } while (superClassEntityAnnotation == null);
      if (superClassEntityAnnotation.name().isEmpty()) {
        kind = superClassElement.getSimpleName().toString();
      } else {
        kind = superClassEntityAnnotation.name();
      }
    } else {
      this.env.getMessager().printMessage(
          Kind.NOTE, 
          "ERROR.: Neither @Entity nor @Subclass annotations were found on '" + entityName + "'.");
      return false;
    }
  
    this.env.getMessager().printMessage(Kind.NOTE, "Kind for '" + entityName + "' is '" + kind + "'.");

    this.env.getMessager().printMessage(Kind.NOTE, "Analyzing fields:");
    String parentKind = null;
    String idType = null;
    for (VariableElement fieldElement : ElementFilter.fieldsIn(env.getElementUtils().getAllMembers(
        (TypeElement) entityElement))) {
        
      Id id = fieldElement.getAnnotation(Id.class);
      if(id != null) {
        idType = fieldElement.asType().toString();
          this.env.getMessager().printMessage(
              Kind.NOTE,  
              "Found @Id " + fieldElement.toString() + " field.");;
      } else {

        Parent parent = fieldElement.getAnnotation(Parent.class);
        if (parent != null) {
          parentKind = env.getTypeUtils().asMemberOf(
              (DeclaredType) entityElement.asType(), fieldElement).toString();
          this.env.getMessager().printMessage(
              Kind.NOTE,  
              "Found @Parent " + fieldElement.toString() + " field.");
        } else {
          this.env.getMessager().printMessage(
              Kind.NOTE,  
              "Ignoring " + fieldElement.toString() + " field.");
        }
      }
    }
    
    if(idType == null) {
      this.env.getMessager().printMessage(Kind.NOTE, "ERROR: Could not find @Id field on '" + entityName + "'.");
      return false;
    }
      

    try {
      out = new PrintWriter(
          new BufferedWriter(this.env.getFiler().createSourceFile(
              keyFactoryQName, entityElement).openWriter()));

      out.println("package " + entityPackageName + ";");
      out.println();
      out.println("import com.google.appengine.api.datastore.KeyFactory;");
      out.println("import com.googlecode.objectify.Key;");
      out.println();
      out.println("public class " + keyFactoryName + " { ");
      out.println();
      if (parentKind == null) {
        out.println("  public static Key<" + entityName
            + "> createKey(" + idType + " id) {");
        out.println("    return new Key<" + entityName
            + ">(KeyFactory.createKey(null, \"" + kind + "\", id));");
      } else {
        out.println("  public static Key<" + entityName + "> createKey("
            + parentKind + " parentKey, " + idType + " id) {");
        out.println("    return new Key<" + entityName
            + ">(KeyFactory.createKey(parentKey.getRaw(), \"" + kind + "\", id));");
      }
      out.println("  }");
      out.println("}");
      out.println();

    } catch (java.io.IOException e) {
      this.env.getMessager().printMessage(Kind.ERROR, e.toString());
      throw new RuntimeException(e);
    } finally {
      if (out != null) {
        out.close();
      }
    }
    return true;
  }
}
